#!/usr/bin/env bash
set -euo pipefail

dn=$(dirname "$0")
# shellcheck source=src/cmdlib.sh
. "${dn}"/cmdlib.sh

# IBM SecureExecution
secure_execution=
image_suffix=

# This script is used for creating both the bare metal and the canonical VM
# image (qemu). `buildextend-qemu` is a symlink to `buildextend-metal`.
case "$(basename "$0")" in
    "cmd-buildextend-metal") image_type=metal;;
    "cmd-buildextend-metal4k") image_type=metal4k;;
    "cmd-buildextend-dasd") image_type=dasd;;
    "cmd-buildextend-qemu") image_type=qemu;;
    "cmd-buildextend-secex")
        secure_execution=1
        image_type=qemu
        image_suffix=-secex
        ;;
    *) fatal "called as unexpected name $0";;
esac

print_help() {
    cat 1>&2 <<EOF
Usage: nestos-assembler buildextend-${image_type} --help
       nestos-assembler buildextend-${image_type} [--build ID]

  Build a bare metal image.
EOF
}

# Parse options
hostkey=
genprotimgvm=/data.secex/genprotimgvm.qcow2
ignition_pubkey=
rc=0
build=
force=
options=$(getopt --options h --longoptions help,force,build:,hostkey:,genprotimgvm: -- "$@") || rc=$?
[ $rc -eq 0 ] || {
    print_help
    exit 1
}
eval set -- "$options"
while true; do
    case "$1" in
        -h | --help)
            print_help
            exit 0
            ;;
        --force)
            force=1
            ;;
        --build)
            build=$2
            shift
            ;;
        --hostkey)
            hostkey=$(realpath "$2")
            shift
            ;;
        --genprotimgvm)
            genprotimgvm="$2"
            shift
            ;;
        --)
            shift
            break
            ;;
        -*)
            fatal "$0: unrecognized option: $1"
            ;;
        *)
            break
            ;;
    esac
    shift
done

if [ $# -ne 0 ]; then
    print_help
    fatal "Too many arguments passed"
fi

case "$basearch" in
    "x86_64"|"aarch64"|"s390x"|"ppc64le"|"riscv64") ;;
    *) fatal "$basearch is not supported for this command" ;;
esac

if [[ "$basearch" != "s390x" && $image_type == dasd ]]; then
    fatal "$basearch is not supported for building dasd images"
fi

# shellcheck disable=SC2031
export LIBGUESTFS_BACKEND=direct
export IMAGE_TYPE="${image_type}"
prepare_build

if [ -z "${build}" ]; then
    build=$(get_latest_build)
    if [ -z "${build}" ]; then
        fatal "No build found."
    fi
fi

builddir=$(get_build_dir "$build")
if [ ! -d "${builddir}" ]; then
    fatal "Build dir ${builddir} does not exist."
fi

# add building sempahore
build_semaphore="${builddir}/.${image_type}.building"
if [ -e "${build_semaphore}" ]; then
    fatal "${build_semaphore} found: another process is building ${image_type}"
fi
touch "${build_semaphore}"
trap 'rm -rf ${build_semaphore}' EXIT

# check if the image already exists in the meta.json
if [ -z "${force}" ]; then
    meta_img=$(meta_key "images.${image_type}${image_suffix}.path")
    if [ "${meta_img}" != "None" ]; then
        echo "${image_type}${image_suffix} image already exists:"
        echo "$meta_img"
        exit 0
    fi
fi

# reread these values from the build itself rather than rely on the ones loaded
# by prepare_build since the config might've changed since then
name=$(meta_key name)
ref=$(meta_key ref)
if [ "${ref}" = "None" ]; then
    ref=""
fi
commit=$(meta_key ostree-commit)

ostree_repo=${tmprepo}
# Ensure that we have the cached unpacked commit
import_ostree_commit_for_build "${build}"
# Note this overwrote the bits generated in prepare_build
# for image_json.  In the future we expect to split prepare_build
# into prepare_ostree_build and prepare_diskimage_build; the
# latter path would only run this.
image_json=${workdir}/tmp/image.json

image_format=raw
if [[ $image_type == qemu ]]; then
    image_format=qcow2
fi

img=${name}-${build}-${image_type}${image_suffix}.${basearch}.${image_format}
path=${PWD}/${img}

ignition_platform_id="${image_type}"

# dasd and metal4k are different disk formats, but they're still metal
if [ "${image_type}" = dasd ] || [ "${image_type}" = metal4k ]; then
    ignition_platform_id=metal
fi

# We do some extra handling of the rootfs here; it feeds into size estimation.
rootfs_type=$(jq -re .rootfs < "${image_json}")

deploy_via_container=""
if jq -re '.["deploy-via-container"]' < "${image_json}"; then
    deploy_via_container="true"
fi
container_imgref=$(jq -r '.["container-imgref"]//""' < "${image_json}")
# Nowadays we pull the container across virtiofs rather than accessing the repo, see
# https://github.com/openshift/os/issues/594
ostree_container=ostree-unverified-image:oci-archive:$builddir/$(meta_key images.ostree.path)

# fs-verity requires block size = page size. We need to take that into account
# in the disk size estimation due to higher fragmentation on larger blocks.
BLKSIZE=""
if [ "${rootfs_type}" = "ext4verity" ]; then
    BLKSIZE="$(getconf PAGE_SIZE)"
fi

disk_args=()
qemu_args=()
# SecureExecution extra stuff
if [[ $secure_execution -eq "1" ]]; then
    ignition_pubkey=$(mktemp -p "${tmp_builddir}")
    disk_args+=("--with-secure-execution" "--write-ignition-pubkey-to" "${ignition_pubkey}")
    if [ -z "${hostkey}" ]; then
        if [ ! -f "${genprotimgvm}" ]; then
            fatal "No genprotimgvm provided at ${genprotimgvm}"
        fi
        genprotimg_img="${PWD}/secex-genprotimg.img"
        qemu-img create -f raw "${genprotimg_img}" 512M
        mkfs.ext4 "${genprotimg_img}"
        qemu_args+=("-drive" "if=none,id=genprotimg,format=raw,file=${genprotimg_img}" \
                            "-device" "virtio-blk,serial=genprotimg,drive=genprotimg")
    else
        qemu_args+=("-drive" "if=none,id=hostkey,format=raw,file=$hostkey,readonly=on" \
                            "-device" "virtio-blk,serial=hostkey,drive=hostkey")
    fi
fi

echo "Estimating disk size..."
# The additional 35% here is obviously a hack, but we can't easily completely fill the filesystem,
# and doing so has apparently negative performance implications.
/usr/lib/coreos-assembler/estimate-commit-disk-size ${BLKSIZE:+--blksize ${BLKSIZE}} --repo "$ostree_repo" "$commit" --add-percent 35 > "$PWD/tmp/ostree-size.json"
rootfs_size="$(jq '."estimate-mb".final' "$PWD/tmp/ostree-size.json")"
# extra size is the non-ostree partitions, see create_disk.sh
image_size="$(( rootfs_size + 513 ))M"
echo "Disk size estimated to ${image_size}"

# For bare metal and dasd images, we use the estimated image size. For IaaS/virt, we get it from
# image.yaml because we want a "default" disk size that has some free space.
case "${image_type}" in
    metal*|dasd)
        # Unset the root size, which will inherit from the image size
        rootfs_size=0
        ;;
    qemu)
        image_size="$(jq -r ".size" < "${image_json}")G"
        rootfs_size="${rootfs_size}M"
        ;;
    *) fatal "unreachable image_type ${image_type}";;
esac

if [ "${image_type}" == metal4k ]; then
    disk_args+=("--no-x86-bios-bootloader")
fi

set -x
kargs="$(python3 -c 'import sys, json; args = json.load(sys.stdin)["extra-kargs"]; print(" ".join(args))' < "${image_json}")"
kargs="$kargs ignition.platform.id=$ignition_platform_id"

qemu-img create -f ${image_format} "${path}.tmp" "${image_size}"

extra_target_device_opts=""
# we need 4096 block size for ECKD DASD and (obviously) metal4k
if [[ $image_type == dasd || $image_type == metal4k ]]; then
  extra_target_device_opts=",physical_block_size=4096,logical_block_size=4096"
fi
qemu_args+=("-drive" "if=none,id=target,format=${image_format},file=${path}.tmp,cache=unsafe" \
              "-device" "virtio-blk,serial=target,drive=target${extra_target_device_opts}")

# Generate the JSON describing the disk we want to build
cat >image-dynamic.json << EOF
{
    "rootfs-size": "${rootfs_size}",
    "osname": "${name}",
    "buildid": "${build}",
    "imgid": "${img}",
    "deploy-via-container": "${deploy_via_container}",
    "container-imgref": "${container_imgref}",
    "ostree-commit": "${commit}",
    "ostree-ref": "${ref}",
    "ostree-container": "${ostree_container}"
}
EOF
cat "${image_json}" image-dynamic.json | jq -s add > image-for-disk.json
platforms_json="${workdir}/tmp/platforms.json"
yaml2json "${configdir}/platforms.yaml" "${platforms_json}"

if [ "${image_type}" == "qemu" ] && [ "${COSA_USE_OSBUILD:-}" != "" ]; then
    runvm -- /usr/lib/coreos-assembler/runvm-osbuild              \
                "${ostree_repo}" "${ref}"                         \
                /usr/lib/coreos-assembler/coreos.osbuild.mpp.yaml \
                "${path}.tmp"
else
    runvm "${qemu_args[@]}" -- \
            /usr/lib/coreos-assembler/create_disk.sh \
                --config "$(pwd)"/image-for-disk.json \
                --kargs "${kargs}" \
                --platform "${ignition_platform_id}" \
                --platforms-json "${platforms_json}" \
                "${disk_args[@]}"
fi

if [[ $secure_execution -eq "1" && -z "${hostkey}" ]]; then
    /usr/lib/coreos-assembler/secex-genprotimgvm-scripts/runvm.sh \
        --genprotimgvm "${genprotimgvm}" -- "${qemu_args[@]}"
fi

/usr/lib/coreos-assembler/finalize-artifact "${path}.tmp" "${path}"

sha256=$(sha256sum_str < "${img}")
cosa meta --workdir "${workdir}" --build "${build}" --dump | python3 -c "
import sys, json
j = json.load(sys.stdin)
j['images']['${image_type}${image_suffix}'] = {
    'path': '${img}',
    'sha256': '${sha256}',
    'size': $(stat -c '%s' "${img}")
}
json.dump(j, sys.stdout, indent=4)
" | jq -s add > "meta.json.new"

# one more artifact for Secure Execution
if [[ -n "${ignition_pubkey}" ]]; then
    gpg_key=${name}-${build}-ignition-secex-key.gpg.pub
    python3 -c "
import sys, json
j = json.load(sys.stdin)
j['images']['ignition-gpg-key'] = {
    'path': '${gpg_key}',
    'sha256': '$(sha256sum_str < "${ignition_pubkey}")',
    'size': $(stat -c '%s' "${ignition_pubkey}"),
    'skip-compression': True
}
json.dump(j, sys.stdout, indent=4)
" < "meta.json.new" | jq -s add > "key.json"
    mv key.json meta.json.new
    /usr/lib/coreos-assembler/finalize-artifact "${ignition_pubkey}" "${builddir}/${gpg_key}"
fi

# and now the crucial bits
cosa meta --workdir "${workdir}" --build "${build}" --artifact "${image_type}" --artifact-json "$(readlink -f meta.json.new)"
/usr/lib/coreos-assembler/finalize-artifact "${img}" "${builddir}/${img}"

# Quiet for the rest of this so the last thing we see is a success message
set +x
# clean up the tmpbuild
rm -rf "${tmp_builddir}"
echo "Successfully generated: ${img}"
